#include "SparseMatrix.hpp"
#include <Time.hpp>
#include <mkl.h>

int main() {
    const int kSize = 100;
    const int kSize2 = kSize+1;
    CSRMatrix a(kSize, kSize2), c(kSize, kSize), d(kSize, kSize);
    CSCMatrix b(kSize2, kSize);
    csr_random_fill(a);
    csc_random_fill(b);

    std::cout << "sp_mm_v1:\t"
              << get_execution_time(std::bind(sp_mm_v1, std::ref(a), std::ref(b), std::ref(d))) << " ms\n";

#define MM_CHECK(csr_csc_mm_vx)                                                                      \
    csr_zero_fill(c);                                                                                \
    csr_csc_mm_vx(a, b, c);                                                                          \
    std::cout << #csr_csc_mm_vx << ":\t";                                                              \
    std::cout << get_execution_time(std::bind(csr_csc_mm_vx, std::ref(a), std::ref(b), std::ref(c))) \
              << " ms";                                                                              \
    if (csr_equal(c, d)) {                                                                           \
        std::cout << "\tpassed \n";                                                                    \
    } else {                                                                                         \
        std::cout << "\tfailed \n";                                                                    \
    }

    MM_CHECK(sp_mm_v2);
    MM_CHECK(sp_mm_v3);
#undef MM_CHECK

    mkl_set_num_threads(1);
    sparse_matrix_t csr_a;
    sparse_matrix_t csc_b;
    sparse_matrix_t csr_d;

    mkl_sparse_d_create_csr(&csr_a, SPARSE_INDEX_BASE_ZERO, a.rowCount_, a.colCount_,
                            a.rowOffset_.data(), a.rowOffset_.data() + 1,
                            a.colInd_.data(),
                            a.data_.data());

    mkl_sparse_d_create_csc(&csc_b, SPARSE_INDEX_BASE_ZERO, b.rowCount_, b.colCount_,
                            b.colOffset_.data(), b.colOffset_.data() + 1,
                            b.rowInd_.data(),
                            b.data_.data());
    mkl_sparse_d_create_csr(&csr_d, SPARSE_INDEX_BASE_ZERO, d.rowCount_, d.colCount_,
                            d.rowOffset_.data(), d.rowOffset_.data() + 1,
                            d.colInd_.data(),
                            d.data_.data());

    std::cout << "sp_mm_mkl:\t"
              << get_execution_time(std::bind(mkl_sparse_spmm, SPARSE_OPERATION_NON_TRANSPOSE, csr_a, csc_b, &csr_d))
              << " ms\n";
    return 0;
}
